import os
import ast
import asttokens
import sys
import astunparse


sys.path.append("..")

from data_structure import FileSystemNode
from neo4jDBOP import neo4jDBOP
from import_graph.source_visitor import FunVisitor, ImportVisitor

neo4j_conf = {
    "uri" : "neo4j://211.71.15.39:7685",
    "user" : "neo4j",
    "password" : "pysr"
}
# browser : 211.71.15.39:7472
# 211.71.15.39:7471

DBOP = neo4jDBOP(neo4j_conf)

class ImportGraph(object):

    def __init__(self, pkg_name, version):
        self.pkg_name = pkg_name
        self.version = version
        self.absolute_path = "/home/sdp/science/sr/PySC_API/pkgs/" + pkg_name + "-" + version # 项目顶层目录 
        self.relative_path = pkg_name + "-" + version
        self.root = None
        self.pyfile_nodes = []
        self.base_dict = {} # {python file: {子类: [父类]}}
        self.import_dict = {} # {python file: {fully_qualified_name:[name, asname, path]}} fully_qualified_name是当前文件+asname，被导入的fully_qualified_name为path+name
        # self.import_count = {} # {python file: [count, flag]}
        self.mapping_flag = False
        self.dp_dict = {} # pakcge的依赖 {package_name:version}

        # first_name = os.path.basename(path)
        # self.root = FileSystemNode(path=path,typ="Folder",striped_name=first_name) # 文件系统的根节点


    def generate_file_system_graph(self, absolute_path, relative_path):
        """ 生成file system graph，为每个python文件生成AST """
        # print(absolute_path, relative_path)
        if os.path.isdir(absolute_path):
            striped_name = os.path.basename(absolute_path)
            fully_qualified_name = relative_path.replace('/', '.')
            node = FileSystemNode(fully_qualified_name=fully_qualified_name, striped_name=striped_name, typ="Folder")
            if self.root is None:
                self.root = node
            init_path = os.path.join(absolute_path, "__init__.py")
            if os.path.exists(init_path):
                DBOP.insertPackage(fully_qualified_name=fully_qualified_name, striped_name=striped_name)
            else:
                DBOP.insertFloder(fully_qualified_name=fully_qualified_name, striped_name=striped_name)
            # print(absolute_path)
            for item in os.listdir(absolute_path):
                child_node = self.generate_file_system_graph(os.path.join(absolute_path, item), os.path.join(relative_path, item))
                if child_node and child_node.typ == "Folder":
                    child_node.parent = node
                    node.children.append(child_node)
                    DBOP.insertFileSystemEdge(node.fully_qualified_name, child_node.fully_qualified_name)
                    # DBOP.insertFileSystemEdgeFather(child_node.fully_qualified_name, node.fully_qualified_name)
                elif child_node and child_node.typ == "Module":
                    child_node.parent = node
                    node.children.append(child_node)
                    DBOP.insertModuleContains(father_fully_qualified_name=node.fully_qualified_name, module_fully_qualified_name=child_node.fully_qualified_name)
            return node
        
        elif os.path.isfile(absolute_path) and absolute_path.endswith(".py"):
            file_name = os.path.basename(absolute_path)
            striped_name = file_name[:-3]
            fully_qualified_name = relative_path.replace('/', '.')[:-3]
            node = FileSystemNode(fully_qualified_name=fully_qualified_name, striped_name=striped_name, typ="Module")
            DBOP.insertModule(fully_qualified_name=fully_qualified_name, striped_name=striped_name)

            self.pyfile_nodes.append(node)
            if self.root is None:
                self.root = node

            with open(absolute_path,"r") as f:
                code_str = f.read()
            try:
                atok = asttokens.ASTTokens(code_str, parse=True)
                tree = atok.tree
            except Exception:
                DBOP.labelModuleASTerror(fully_qualified_name=fully_qualified_name, striped_name=striped_name)
                return
            # print(astunparse.dump(tree))
            
            node.ast_tree = tree
            visitor = FunVisitor()
            visitor.visit(tree)
            # print(visitor.pair)
            if len(visitor.pair) > 0:
                self.base_dict[fully_qualified_name] = visitor.pair
            # print(visitor.result)

            for striped_name in visitor.result:
                item = visitor.result[striped_name]
                if isinstance(item, tuple): # function
                    function_fully_qualified_name = fully_qualified_name + '.' + striped_name
                    DBOP.insertFunction(fully_qualified_name=function_fully_qualified_name, striped_name=striped_name, args=item)
                    DBOP.insertFunctionContains(father_fully_qualified_name=fully_qualified_name, function_fully_qualified_name=function_fully_qualified_name)
                elif isinstance(item, dict):
                    class_fully_qualified_name = fully_qualified_name + '.' + striped_name
                    DBOP.insertClass(fully_qualified_name=class_fully_qualified_name, striped_name=striped_name)
                    DBOP.insertClassContains(father_fully_qualified_name=fully_qualified_name, class_fully_qualified_name=class_fully_qualified_name)
                    for fun_name in item:
                        fun_args = item[fun_name]
                        function_fully_qualified_name = class_fully_qualified_name + '.' + fun_name
                        DBOP.insertFunction(fully_qualified_name=function_fully_qualified_name, striped_name=fun_name, args=fun_args)
                        DBOP.insertFunctionContains(father_fully_qualified_name=class_fully_qualified_name, function_fully_qualified_name=function_fully_qualified_name)
                    # if name in visitor.pair and visitor.pair[name] is not None:
                    #     if relative_path_dbop not in self.base_dict:
                    #         self.base_dict[relative_path_dbop] = {}
                    #     self.base_dict[relative_path_dbop][class_fully_qualified_name] = visitor.pair[name]

            return node
        
    def API_mapping(self):
        for node in self.pyfile_nodes:
            # print(node.striped_name)
            visitor = ImportVisitor(self.dp_dict)
            visitor.visit(node.ast_tree)
            # print("*****: ", node.striped_name)
            # print("import_dict:", visitor.import_dict)
            # print("importfrom_dict: ", visitor.importfrom_dict)
            # import_dict = {} # {name:asname}
            # importfrom_dict = {} # {level: {prefix_path: [[name, asname]]}
            # if node.fully_qualified_name not in self.import_count:
            #     self.import_count[node.fully_qualified_name] = [0, False]
            for name in visitor.import_dict:
                if node.fully_qualified_name not in self.import_dict:
                    self.import_dict[node.fully_qualified_name] = {}
                if visitor.import_dict[name] is None:
                    new_name = name
                else:
                    new_name = visitor.import_dict[name]
                fully_qualified_name = node.fully_qualified_name + "." + new_name
                self.import_dict[node.fully_qualified_name][fully_qualified_name] = [name, visitor.import_dict[name], None]
                # self.import_count[node.fully_qualified_name][0] += 1
            for level in visitor.importfrom_dict:
                for path in visitor.importfrom_dict[level]:
                    for name in visitor.importfrom_dict[level][path]:
                        if node.fully_qualified_name not in self.import_dict:
                            self.import_dict[node.fully_qualified_name] = {}
                        if name[1] is None:
                            new_name = name[0]
                        else:
                            new_name = name[1]
                        fully_qualified_name = node.fully_qualified_name + "." + new_name
                        self.import_dict[node.fully_qualified_name][fully_qualified_name] = [name[0], name[1], path]
                        # self.import_count[node.fully_qualified_name][0] += 1
            # print(self.root.striped_name)
        # print(self.import_dict)
        pre_count = DBOP.getNodeCount()
        while True:
            for file_name in self.import_dict:
                    for name in self.import_dict[file_name]:
                        import_info = self.import_dict[file_name][name]
                        if file_name.endswith("__init__"):
                            father_name = file_name[:-9]
                        else:
                            father_name = file_name
                        self.mapping_import(file_fully_qualified_name=father_name, import_info=(import_info[0], import_info[1], import_info[2]))
                            # if flag:
                            #     import_info[3] = True
                                # file_mapping_count += 1
            now_count = DBOP.getNodeCount()
            if pre_count == now_count: break
            pre_count = now_count
        
    def mapping_import(self, file_fully_qualified_name, import_info):
        """ node是当前分析import/import from的python文件节点
            import_info为(name, asname, path)
        # """
        # print()
        # print("file name: ", file_fully_qualified_name)
        # print("import_info: ", import_info)
        if import_info[1] is not None:
            new_name = import_info[1]
        else:
            new_name = import_info[0]
        # print("new_name: ", new_name)
        # if node.fully_qualified_name not in self.import_dict:
        #     self.import_dict[node.fully_qualified_name] = {}
        new_fully_qualified_name = file_fully_qualified_name + "." + new_name

        # print("new_fully_qualified_name: ", new_fully_qualified_name)
        # self.import_dict[node.fully_qualified_name][new_name] = importfrom_info[2] + importfrom_info[0]
        fully_path = ""
        if import_info[2] is not None:
            fully_path += import_info[2] + "."
        fully_path += import_info[0]
        result = DBOP.searchPath(prefix=self.root.striped_name, path=fully_path)
        if len(result) == 0:
            return
        
        if result[0][2] == "Function":
            self.mapping_function(old_fully_qualified_name=result[0][0], new_fully_qualified_name=new_fully_qualified_name, new_name=new_name, father_fully_qualified_name=file_fully_qualified_name)
        elif result[0][2] == "Class":
            self.mapping_class(old_fully_qualified_name=result[0][0], new_fully_qualified_name=new_fully_qualified_name, new_name=new_name, father_fully_qualified_name=file_fully_qualified_name)
        elif result[0][2] == "Module":
            self.mapping_module(old_fully_qualified_name=result[0][0], new_fully_qualified_name=new_fully_qualified_name, new_name=new_name, father_fully_qualified_name=file_fully_qualified_name)
        elif result[0][2] == "Package":
            self.mapping_package(old_fully_qualified_name=result[0][0], new_fully_qualified_name=new_fully_qualified_name, new_name=new_name, father_fully_qualified_name=file_fully_qualified_name)
        return
        
    
    def mapping_function(self, old_fully_qualified_name, new_fully_qualified_name, new_name, father_fully_qualified_name):
        """ node节点表示新添加function的文件
            old_fully_qualified_name表示import mapping的目标节点
            new_fully_qualified_name表示import mapping的源节点，即新添加的节点
        """
        # new_fully_qualified_name = node.fully_qualified_name + "." + new_name
        fun_info = DBOP.searchFunctionInfo(fully_qualified_name=old_fully_qualified_name)
        DBOP.insertFunction(fully_qualified_name=new_fully_qualified_name, striped_name=new_name, args=fun_info[0])
        DBOP.insertFunctionContains(father_fully_qualified_name=father_fully_qualified_name, function_fully_qualified_name=new_fully_qualified_name)
        DBOP.mappingImport(old_name=old_fully_qualified_name, new_name=new_fully_qualified_name)

    def mapping_class(self, old_fully_qualified_name, new_fully_qualified_name, new_name, father_fully_qualified_name):
        """ node节点表示新添加class的文件
            old_fully_qualified_name表示import mapping的目标节点
            new_fully_qualified_name表示import mapping的源节点，即新添加的节点
        """
        # new_fully_qualified_name = node.fully_qualified_name + "." + new_name
        DBOP.insertClass(fully_qualified_name=new_fully_qualified_name, striped_name=new_name)
        DBOP.insertClassContains(father_fully_qualified_name=father_fully_qualified_name, class_fully_qualified_name=new_fully_qualified_name)
        DBOP.mappingImport(old_name=old_fully_qualified_name, new_name=new_fully_qualified_name)
        fun_list = DBOP.searchFunction(fully_qualified_name=old_fully_qualified_name)
        for fun_name in fun_list:
            new_function_name = new_fully_qualified_name + "." + fun_name[1]
            self.mapping_function(old_fully_qualified_name=fun_name[0], new_fully_qualified_name=new_function_name, new_name=fun_name[1], father_fully_qualified_name=new_fully_qualified_name)
            # fun_info = DBOP.searchPath(prefix=self.root.fully_qualified_name, path=fun_name[0])
            # DBOP.insertFunction(fully_qualified_name=new_function_name, striped_name=fun_name[1], args=fun_info[0])
            # DBOP.insertFunctionContains(father_fully_qualified_name=new_fully_qualified_name, function_fully_qualified_name=new_function_name)
            # DBOP.mappingImport(old_name=fun_name[0], new_name=new_function_name)

    def mapping_module(self, old_fully_qualified_name, new_fully_qualified_name, new_name, father_fully_qualified_name):
        """ node节点表示新添加module的文件
            old_fully_qualified_name表示import mapping的目标节点
            new_fully_qualified_name表示import mapping的源节点，即新添加的节点
        """
        # new_fully_qualified_name = node.fully_qualified_name + "." + new_name
        DBOP.insertModule(fully_qualified_name=new_fully_qualified_name, striped_name=new_name)
        DBOP.insertModuleContains(father_fully_qualified_name=father_fully_qualified_name, module_fully_qualified_name=new_fully_qualified_name)
        DBOP.mappingImport(old_name=old_fully_qualified_name, new_name=new_fully_qualified_name)
        package_list = DBOP.searchPackage(fully_qualified_name=old_fully_qualified_name)
        # print("package_list: ", package_list)
        for package_name in package_list:
            new_package_fully_qualified_name = new_fully_qualified_name + "." + package_name[1]
            self.mapping_package(old_fully_qualified_name=package_name[1], new_fully_qualified_name=new_package_fully_qualified_name, new_name=package_name[1], father_fully_qualified_name=new_fully_qualified_name)
        
        class_list = DBOP.searchClass(fully_qualified_name=old_fully_qualified_name)
        # print("class_list: ", class_list)
        for class_name in class_list:
            new_class_name = new_fully_qualified_name + "." + class_name[1]
            self.mapping_class(old_fully_qualified_name=class_name[0], new_fully_qualified_name=new_class_name, new_name=class_name[1], father_fully_qualified_name=new_fully_qualified_name)
            # DBOP.insertClass(fully_qualified_name=new_class_name, striped_name=class_name[1])
            # DBOP.insertClassContains(father_fully_qualified_name=new_fully_qualified_name, class_fully_qualified_name=new_class_name)
            # DBOP.mappingImport(old_name=class_name[0], new_name=new_class_name)
            # fun_list = DBOP.searchClassFunction(class_name=class_name[0])
            # for fun_name in fun_list:
            #     new_function_name = new_class_name + "." + fun_name[1]
            #     fun_info = DBOP.searchPath(prefix=self.root.fully_qualified_name, path=fun_name[0])
            #     DBOP.insertFunction(fully_qualified_name=new_function_name, striped_name=fun_name[1], args=fun_info[0])
            #     DBOP.insertFunctionContains(father_fully_qualified_name=new_class_name, function_fully_qualified_name=new_function_name)
            #     DBOP.mappingImport(old_name=fun_name[0], new_name=new_function_name)
        
        fun_list = DBOP.searchFunction(fully_qualified_name=old_fully_qualified_name)
        for fun_name in fun_list:
            new_function_name = new_fully_qualified_name + "." + fun_name[1]
            self.mapping_function(old_fully_qualified_name=fun_name[0], new_fully_qualified_name=new_function_name, new_name=fun_name[1], father_fully_qualified_name=new_fully_qualified_name)
            # fun_info = DBOP.searchPath(prefix=self.root.fully_qualified_name, path=fun_name[0])
            # DBOP.insertFunction(fully_qualified_name=new_function_name, striped_name=fun_name[1], args=fun_info[0])
            # DBOP.insertFunctionContains(father_fully_qualified_name=new_fully_qualified_name, function_fully_qualified_name=new_function_name)
            # DBOP.mappingImport(old_name=fun_name[0], new_name=new_function_name)
    
    def mapping_package(self, old_fully_qualified_name, new_fully_qualified_name, new_name, father_fully_qualified_name):
        """ node节点表示新添加package的文件
            old_fully_qualified_name表示import mapping的目标节点
            new_fully_qualified_name表示import mapping的源节点，即新添加的节点
        """
        # new_fully_qualified_name = node.fully_qualified_name + "." + new_name
        DBOP.insertPackage(fully_qualified_name=new_fully_qualified_name, striped_name=new_name)
        DBOP.insertPackageContains(father_fully_qualified_name=father_fully_qualified_name, package_fully_qualified_name=new_fully_qualified_name)
        DBOP.mappingImport(old_name=old_fully_qualified_name, new_name=new_fully_qualified_name)
        package_list = DBOP.searchPackage(fully_qualified_name=old_fully_qualified_name)
        # print("package_list: ", package_list)
        for package_name in package_list:
            new_package_fully_qualified_name = new_fully_qualified_name + "." + package_name[1]
            self.mapping_package(old_fully_qualified_name=package_name[1], new_fully_qualified_name=new_package_fully_qualified_name, new_name=package_name[1])
        module_list = DBOP.searchModule(fully_qualified_name=old_fully_qualified_name)
        # print("module_list: ", module_list)
        for module_name in module_list:
            new_module_name = new_fully_qualified_name + "." + module_name[1]
            self.mapping_module(old_fully_qualified_name=module_name[0], new_fully_qualified_name=new_module_name, new_name=module_name[1], father_fully_qualified_name=new_fully_qualified_name)

        fun_list = DBOP.searchFunction(fully_qualified_name=old_fully_qualified_name)
        for fun_name in fun_list:
            new_function_name = new_fully_qualified_name + "." + fun_name[1]
            self.mapping_function(old_fully_qualified_name=fun_name[0], new_fully_qualified_name=new_function_name, new_name=fun_name[1], father_fully_qualified_name=new_fully_qualified_name)
           